from selenium import webdriver
from selenium.webdriver.edge.service import Service
from selenium.webdriver.edge.options import Options
import requests
import time
import os
import json
from datetime import datetime, timedelta
from PIL import Image, ImageDraw
from urllib.parse import quote

# 目标作者名称
TARGET_AUTHOR = "康复～数智掘金"

def get_api_data(url, headers, retries=3, delay=5):
    """获取 API 数据，带重试机制"""
    for attempt in range(retries):
        try:
            response = requests.get(url, headers=headers, timeout=30)
            if response.status_code == 200:
                json_data = response.json()
                if 'resp_data' in json_data and 'topics' in json_data['resp_data']:
                    return json_data
                else:
                    print("数据格式不正确，未包含 'resp_data' 或 'topics' 字段，正在重试...")
            else:
                print(f"请求失败，状态码: {response.status_code}，正在重试...")
        except Exception as e:
            print(f"请求异常: {e}，正在重试...")
        
        time.sleep(delay)
    print("所有重试均失败。")
    return None

def format_create_time(create_time_str):
    """格式化时间：将 create_time 格式化为 'YYYY-MM-DDTHH:MM:SS'"""
    try:
        dt = datetime.strptime(create_time_str.split(".")[0], '%Y-%m-%dT%H:%M:%S')
        return dt.strftime('%Y-%m-%dT%H:%M:%S')
    except ValueError as e:
        print(f"时间格式转换失败: {e}")
        return create_time_str

def get_date_from_time(time_str):
    """从时间字符串中提取日期"""
    try:
        return time_str.split("T")[0]
    except:
        return None

def download_and_process_avatar(avatar_url, avatar_path, author):
    """下载并处理头像为圆形"""
    try:
        if os.path.exists(avatar_path):
            print(f"头像 {os.path.basename(avatar_path)} 已存在，跳过下载。")
            return True
            
        avatar_data = requests.get(avatar_url, timeout=30).content
        with open(avatar_path, 'wb') as avatar_file:
            avatar_file.write(avatar_data)

        avatar_image = Image.open(avatar_path).convert("RGBA")
        avatar_image = avatar_image.resize((300, 300), Image.LANCZOS)

        # 创建圆形头像蒙版
        mask = Image.new('L', (300, 300), 0)
        draw = ImageDraw.Draw(mask)
        draw.ellipse((0, 0, 300, 300), fill=255)

        # 将蒙版应用到头像
        avatar_image.putalpha(mask)
        avatar_image.save(avatar_path)
        print(f"圆形头像 {os.path.basename(avatar_path)} 已成功保存。")
        return True
    except Exception as e:
        print(f"处理头像 {os.path.basename(avatar_path)} 时出错: {e}")
        return False

def download_image(image_url, image_path):
    """下载图片"""
    try:
        if os.path.exists(image_path):
            return True
        img_data = requests.get(image_url, timeout=30).content
        with open(image_path, 'wb') as img_file:
            img_file.write(img_data)
        print(f"图片 {os.path.basename(image_path)} 已成功下载。")
        return True
    except Exception as e:
        print(f"下载图片 {os.path.basename(image_path)} 时出错: {e}")
        return False

def get_topics_for_date(date_str, headers, group_id, count=20):
    """获取指定日期当天的所有文章（分页处理）"""
    all_topics = []
    end_time = f"{date_str}T23:59:59.999+0800"
    end_time_encoded = quote(end_time, safe='')
    
    page = 0
    while True:
        url = f"https://api.zsxq.com/v2/groups/{group_id}/topics?scope=digests&count={count}&end_time={end_time_encoded}"
        print(f"正在获取 {date_str} 的数据，第 {page + 1} 页...")
        
        json_data = get_api_data(url, headers)
        if not json_data:
            break
            
        topics = json_data['resp_data'].get('topics', [])
        if not topics:
            break
        
        # 过滤出目标作者的内容，并检查日期范围
        filtered_topics = []
        found_earlier_date = False
        
        for topic in topics:
            try:
                topic_date = get_date_from_time(topic.get('create_time', ''))
                
                # 如果文章日期早于目标日期，说明已经超出范围，停止分页
                if topic_date and topic_date < date_str:
                    found_earlier_date = True
                    break
                
                # 只处理目标日期和目标作者的内容
                if topic_date == date_str:
                    author = topic['talk']['owner']['name']
                    if author == TARGET_AUTHOR:
                        filtered_topics.append(topic)
            except:
                continue
        
        all_topics.extend(filtered_topics)
        
        # 如果已经超出日期范围，停止分页
        if found_earlier_date:
            break
        
        # 检查是否还有更多数据
        if len(topics) < count:
            break
        
        # 更新 end_time 为当前批次最早的时间
        last_topic = topics[-1]
        if 'create_time' in last_topic:
            last_topic_date = get_date_from_time(last_topic['create_time'])
            # 如果最后一篇文章的日期已经早于目标日期，停止分页
            if last_topic_date and last_topic_date < date_str:
                break
            end_time = last_topic['create_time']
            end_time_encoded = quote(end_time, safe='')
        else:
            break
        
        page += 1
        time.sleep(2)  # 避免请求过快
    
    print(f"日期 {date_str} 共获取到 {len(all_topics)} 篇目标作者的文章")
    return all_topics

def process_topics_to_markdown(topics_by_date, avatar_path, author_name, author_info):
    """将按日期分组的话题转换为 markdown"""
    markdown_content = []
    
    # 在文档开头显示用户信息（只显示一次）
    markdown_content.append(f"![头像]({avatar_path})\n\n")
    markdown_content.append(f"**{author_name}**\n\n")
    if author_info:
        markdown_content.append(f"{author_info}\n\n")
    markdown_content.append("---\n\n")
    
    # 按日期排序（从早到晚）
    sorted_dates = sorted(topics_by_date.keys())
    
    for date_str in sorted_dates:
        topics = topics_by_date[date_str]
        if not topics:
            continue
        
        # 日期分隔
        markdown_content.append(f"\n{'='*60}\n")
        markdown_content.append(f"## {date_str}\n\n")
        
        # 处理该日期的所有文章
        for topic in topics:
            try:
                talk = topic['talk']
                text = talk.get('text', '')
                create_time = topic.get('create_time', '')
                
                if create_time:
                    create_time = format_create_time(create_time)
                    time_part = create_time.split('T')[1] if 'T' in create_time else ''
                else:
                    time_part = "未知时间"
                
                # 添加创建时间
                markdown_content.append(f"### {time_part}\n\n")
                
                # 处理图片（多图遍历）
                images = talk.get('images', [])
                if images:
                    for image in images:
                        try:
                            large_image_url = image['large']['url']
                            image_id = str(image['image_id'])
                            image_filename = f"{image_id}.jpg"
                            image_path = os.path.join("downloaded_images", image_filename)
                            
                            if download_image(large_image_url, image_path):
                                markdown_content.append(f"![图片]({image_path})\n\n")
                        except Exception as e:
                            print(f"处理图片时出错: {e}")
                
                # 添加内容
                if text:
                    markdown_content.append(f"{text}\n\n")
                
                markdown_content.append("---\n\n")
                
            except Exception as e:
                print(f"处理话题时出错: {e}")
                continue
    
    return "".join(markdown_content)

# 主程序
if __name__ == "__main__":
    # 使用 Selenium 启动 Edge 浏览器获取 Cookie
    driver_path = os.path.join(os.path.dirname(__file__), "edgedriver_win64", "msedgedriver.exe")
    if os.path.exists(driver_path):
        service = Service(driver_path)
        driver = webdriver.Edge(service=service)
        print(f"使用本地驱动: {driver_path}")
    else:
        driver = webdriver.Edge()
        print("使用系统 PATH 中的 Edge 驱动")

    driver.get('https://wx.zsxq.com/')
    print("请使用微信扫码登录知识星球...")
    time.sleep(10)  # 等待扫码
    cookies = driver.get_cookies()
    cookie_string = "; ".join([f"{cookie['name']}={cookie['value']}" for cookie in cookies])
    print(f"获取到的 Cookie: {cookie_string[:50]}...")
    driver.quit()

    # 设置目录
    image_dir = "downloaded_images"
    avatar_dir = "avatars"
    os.makedirs(image_dir, exist_ok=True)
    os.makedirs(avatar_dir, exist_ok=True)

    headers = {'Cookie': cookie_string}
    group_id = "458115241428"
    
    # 从 2018-11-01 开始到今天
    start_date = datetime(2018, 11, 1)
    end_date = datetime.now()
    current_date = start_date
    
    # 存储按日期分组的话题
    topics_by_date = {}
    author_info = None
    author_avatar_path = None
    
    print(f"开始从 {start_date.strftime('%Y-%m-%d')} 到 {end_date.strftime('%Y-%m-%d')} 轮询数据...")
    
    # 按日期轮询
    while current_date <= end_date:
        date_str = current_date.strftime('%Y-%m-%d')
        print(f"\n正在处理日期: {date_str}")
        
        # 获取该日期的所有文章
        topics = get_topics_for_date(date_str, headers, group_id, count=20)
        
        if topics:
            topics_by_date[date_str] = topics
            
            # 获取作者信息（从第一篇获取）
            if not author_info and topics:
                try:
                    first_topic = topics[0]
                    owner = first_topic['talk']['owner']
                    author_name = owner['name']
                    avatar_url = owner['avatar_url']
                    
                    # 下载头像
                    avatar_filename = f"{author_name}.png".replace(" ", "_").replace("～", "_")
                    author_avatar_path = os.path.join(avatar_dir, avatar_filename)
                    download_and_process_avatar(avatar_url, author_avatar_path, author_name)
                    
                    # 获取作者介绍（如果有的话）
                    author_info = owner.get('description', '') or owner.get('bio', '') or ''
                    print(f"作者: {author_name}")
                except Exception as e:
                    print(f"获取作者信息时出错: {e}")
        
        # 移动到下一天
        current_date += timedelta(days=1)
        
        # 避免请求过快
        time.sleep(1)
    
    # 生成 markdown
    if topics_by_date:
        print("\n开始生成 Markdown 文档...")
        author_name = TARGET_AUTHOR
        if not author_avatar_path:
            # 如果没有获取到头像，使用默认路径
            avatar_filename = f"{author_name}.png".replace(" ", "_").replace("～", "_")
            author_avatar_path = os.path.join(avatar_dir, avatar_filename)
        
        markdown_content = process_topics_to_markdown(
            topics_by_date, 
            author_avatar_path, 
            author_name,
            author_info or ""
        )
        
        # 保存文件
        filename = f"{TARGET_AUTHOR}_全部内容.md".replace(" ", "_").replace("～", "_")
        with open(filename, 'w', encoding='utf-8') as f:
            f.write(markdown_content)
        print(f"\nMarkdown 内容已成功导出为 '{filename}'")
        print(f"共处理 {len(topics_by_date)} 个日期，总计文章数: {sum(len(topics) for topics in topics_by_date.values())}")
    else:
        print("未找到任何目标作者的内容。")
